package com.freedy.dataStructure.tree;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.Map.Entry;
import java.util.stream.Collectors;

/**
 * @author Freedy
 * @date 2021/3/25 22:31
 */
public class HuffmanCode {
    public static void main(String[] args) {
        zipFile("C:\\Users\\Freedy\\Desktop\\资料\\java数据结构.md",
                "C:\\Users\\Freedy\\Desktop\\资料\\obj.txt");
        unzipFile("C:\\Users\\Freedy\\Desktop\\资料\\obj.txt",
                "C:\\Users\\Freedy\\Desktop\\资料\\fin.md");
    }

    public static void unzipFile(String srcFile,String desFile){
        try(FileInputStream is = new FileInputStream(srcFile);
            ObjectInputStream ooi = new ObjectInputStream(is);
            FileOutputStream fos = new FileOutputStream(desFile))
        {
            final byte[] bytes = decodeFile((HuffmanData) ooi.readObject());
            fos.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void zipFile(String srcFile,String desFile){
        try(FileInputStream is = new FileInputStream(srcFile);
            FileOutputStream os = new FileOutputStream(desFile);
            ObjectOutputStream oos = new ObjectOutputStream(os))
        {
            byte[] b=new byte[is.available()];
            is.read(b);
            HuffmanData encode = encode(b);
            oos.writeObject(encode);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static HuffmanData encode(String str) {
        byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
        return HuffmanEncode(
                bytes, generateCode(
                        buildHuffmanTree(statisticsNodes(
                                bytes
                        ))
                ));
    }

    public static HuffmanData encode(byte[] bytes) {
        return HuffmanEncode(
                bytes, generateCode(
                        buildHuffmanTree(statisticsNodes(
                                bytes
                        ))
                ));
    }

    public static String decode(HuffmanData data) {
        return new String(HuffmanDecode(data.data,data.map,data.flag));
    }

    public static byte[] decodeFile(HuffmanData data) {
        return HuffmanDecode(data.data,data.map,data.flag);
    }

    /**
     * 根据bytes数组统计各个单词出现的频率
     */
    public static List<BinaryTree.Node<HuffmanStruct>> statisticsNodes(byte[] bytes) {
        HashMap<Byte, Integer> map = new HashMap<>();
        for (Byte b : bytes) {
            map.merge(b, 1, Integer::sum);
        }
        return map.entrySet().
                stream().map(item -> new BinaryTree.Node<>(
                new HuffmanStruct(item.getKey(), item.getValue())
        )).sorted((a, b) -> b.getData().getWeight() - a.getData().getWeight()).collect(Collectors.toList());
    }

    /**
     * 构建哈夫曼树
     */
    private static BinaryTree.Node<HuffmanStruct> buildHuffmanTree(List<BinaryTree.Node<HuffmanStruct>> nodes) {
        while (nodes.size() > 1) {
            nodes.sort(Comparator.comparingInt(a -> a.getData().getWeight()));
            BinaryTree.Node<HuffmanStruct> left = nodes.get(0);
            BinaryTree.Node<HuffmanStruct> right = nodes.get(1);
            HuffmanStruct struct = new HuffmanStruct(left.getData().getWeight() + right.getData().getWeight());
            BinaryTree.Node<HuffmanStruct> parent = new BinaryTree.Node<HuffmanStruct>(struct, left, right);
            nodes.remove(0);
            nodes.remove(0);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    /**
     * 生成各个单词对应的编码
     */
    public static Map<Byte, String> generateCode(BinaryTree.Node<HuffmanStruct> node) {
        Map<Byte, String> map = new HashMap<>();
        StringBuilder route = new StringBuilder();
        Stack<BinaryTree.Node<HuffmanStruct>> stack = new Stack<>();
        int deep = 0;
        while (!stack.isEmpty() || node != null) {
            while (node != null) {
                node.getData().setDeep(deep);
                stack.push(node);
                node = node.getLeftNode();
                deep++;
                route.insert(deep - 1, "0");
            }
            if (!stack.isEmpty()) {
                BinaryTree.Node<HuffmanStruct> sNode = stack.pop();
                if (sNode.getRightNode() == null && sNode.getLeftNode() == null) {
                    map.put(sNode.getData().getData(), route.substring(0, sNode.getData().getDeep()));
                }
                node = sNode.getRightNode();
                deep = sNode.getData().getDeep() + 1;
                route.insert(deep - 1, "1");
            }
        }
        return map;
    }

    /**
     * 生成哈夫曼编码
     */
    public static HuffmanData HuffmanEncode(byte[] data, Map<Byte, String> map) {
        StringBuilder builder = new StringBuilder();
        for (byte datum : data) {
            String code = map.get(datum);
            builder.append(code);
        }
        int len = (builder.length() + 7) / 8;
        byte[] encodeByte = new byte[len];
        boolean flag = true;
        for (int i = 0, index = 0; i < builder.length(); i += 8, index++) {
            String stringByte;
            if (i + 8 > builder.length()) {
                stringByte = builder.substring(i);
                flag = false;
            } else {
                stringByte = builder.substring(i, i + 8);
            }
            encodeByte[index] = (byte) Integer.parseInt(stringByte, 2);
        }
        HuffmanData huffmanData = new HuffmanData();
        huffmanData.setData(encodeByte);
        huffmanData.setFlag(flag);
        huffmanData.setMap(map);
        return huffmanData;
    }

    /**
     * 哈夫曼解码
     */
    public static byte[] HuffmanDecode(byte[] decodeByte, Map<Byte, String> map,boolean flag) {
        StringBuilder str = new StringBuilder();
        for (int i = 0; i < decodeByte.length; i++) {
            String s = toBinaryString(flag||(i != decodeByte.length - 1), decodeByte[i]);
            str.append(s);
        }

        StringBuilder builder = new StringBuilder();
        ArrayList<Byte> result = new ArrayList<>();
        for (int i = 0; i < str.length(); i++) {
            builder.append(str.charAt(i));
            if (map.containsValue(builder.toString())) {
                Set<Byte> bytes = map.keySet();
                for (Byte aByte : bytes) {
                    if (map.get(aByte).equals(builder.toString())) {
                        result.add(aByte);
                        break;
                    }
                }
                builder = new StringBuilder();
            }
        }
        byte[] bytes = new byte[result.size()];
        for (int i = 0; i < result.size(); i++) {
            bytes[i] = result.get(i);
        }
        return bytes;
    }

    /**
     *字节转二进制字符串
     */
    private static String toBinaryString(boolean flag, byte b) {
        int temp = b;
        if (flag) {
            temp |= 0x100;
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }


    public static class HuffmanStruct {
        private Byte data;
        private int weight;
        private int deep;

        public int getDeep() {
            return deep;
        }

        public void setDeep(int deep) {
            this.deep = deep;
        }

        public HuffmanStruct(Byte data, int weight) {
            this.data = data;
            this.weight = weight;
        }

        public HuffmanStruct(int weight) {
            this.weight = weight;
        }

        public Byte getData() {
            return data;
        }

        public void setData(Byte data) {
            this.data = data;
        }

        public int getWeight() {
            return weight;
        }

        public void setWeight(int weight) {
            this.weight = weight;
        }

        @Override
        public String toString() {
            return new StringJoiner(", ", HuffmanStruct.class.getSimpleName() + "[", "]")
                    .add("data=" + data)
                    .add("weight=" + weight)
                    .toString();
        }
    }

    public static class HuffmanData implements Serializable{
        private byte[] data;
        private Map<Byte, String> map;
        private boolean flag;

        public HuffmanData(byte[] data, Map<Byte, String> map, boolean flag) {
            this.data = data;
            this.map = map;
            this.flag = flag;
        }


        public HuffmanData() {

        }

        public byte[] getData() {
            return data;
        }

        public void setData(byte[] data) {
            this.data = data;
        }

        public Map<Byte, String> getMap() {
            return map;
        }

        public void setMap(Map<Byte, String> map) {
            this.map = map;
        }

        public boolean isFlag() {
            return flag;
        }

        public void setFlag(boolean flag) {
            this.flag = flag;
        }
    }
}
